# -*- coding:utf-8 -*-
"""
@FileName  :etf_template.py
@Time      :2022/10/26 13:48
@Author    :fsksf

对策略模板进行改造，使其支持多标的买卖、持仓，篮子，申赎
"""
import time
from copy import copy
from typing import Any, Dict
import collections
from vnpy import WORK_DIR
from vnpy_ctastrategy.template import CtaTemplate
from vnpy.trader.utility import round_to
from vnpy.trader.constant import Interval, Direction, Offset, OrderType, Exchange, ProductStatus
from vnpy.trader.object import (
    BarData, TickData, OrderData, TradeData, SubscribeRequest, OrderRequest,
    ContractData, ReportStrategy
)


class ETFTemplate(CtaTemplate):

    author = 'kyq'
    pre_ss_vol = 900000
    per_order_vol = 900000
    parameters = ['pre_ss_vol', 'per_order_vol']         # 申赎最小单位
    trade_basket = True

    def __init__(
        self,
        cta_engine: Any,
        strategy_name: str,
        vt_symbol: str,         # 篮子对应的ETF
        setting: dict,
        trade_basket: bool = True  # 是否交易篮子，关系到是否订阅ETF对应的篮子
    ):
        """"""
        super(ETFTemplate, self).__init__(cta_engine,
                                          strategy_name,
                                          vt_symbol,  # 篮子对应的ETF
                                          setting)
        self.cta_engine = cta_engine
        self.strategy_name = strategy_name
        self.vt_symbol = vt_symbol
        self.trade_basket = trade_basket

        self.inited = False
        self.trading = False
        self.pos = collections.defaultdict(lambda: 0)
        self.target_basket_pos = 0          # 篮子目标数量
        self.require_basket_pos = {}        # 篮子成分股每个股票还需要买多少
        self.basket_pos = 0                 # 篮子包数量（成分股折算）
        self.etf_pos = 0                    # etf的数量(折算)
        self.all_pos = 0                    # 总量 etf + basket_vol （折算）
        self.pre_ss_vol = 900000            # 申赎最小单位
        self.per_order_vol = 900000         # 每次买卖ETF的数量，拆单，每次委托最高一百万(柜台规则)
        self.last_order_time = None

        self.variables = copy(self.variables)
        self.variables.insert(0, "inited")
        self.variables.insert(1, "trading")
        self.variables.insert(2, "etf_pos")
        self.variables.insert(3, "basket_pos")
        self.variables.insert(4, "pos")
        self.variables.insert(6, "all_pos")

        self.update_setting(setting)

    def get_etf_stocks_sub_reqs(self):
        if self.trade_basket:
            return [SubscribeRequest(symbol=comp.symbol,
                                     exchange=comp.exchange,
                                     important=False)
                    for comp in self.cta_engine.main_engine.get_basket_components(self.vt_symbol)]
        else:
            return []

    def on_trade(self, trade: TradeData):
        self.etf_pos = self.pos[self.vt_symbol]
        self.calc_basket_pos()

    def on_start(self):
        self.calc_basket_pos()

    def calc_basket_pos(self):
        """计算持仓"""
        self.basket_pos = self.calculate_basket_pos_by_market(self.pos)
        self.etf_pos = self.pos[self.vt_symbol] / self.pre_ss_vol
        self.all_pos = self.etf_pos + self.basket_pos
        print(f"计算持仓: basket_pos  {self.basket_pos} ,etf_pos  {self.etf_pos}, 持仓 {self.all_pos}")

    def calculate_basket_pos_by_market(self, pos_dict: Dict[str, float]):
        """计算篮子市值"""
        basket_pos = 0
        etf_contract: ContractData = self.cta_engine.main_engine.get_contract(self.vt_symbol)
        stand_basket_market_value = self.cta_engine.main_engine.calculate_stand_basket_market(self.vt_symbol)
        # 如果都等于0，继续向下会卡死
        if stand_basket_market_value == 0:
            return 0
        leg_pos = copy(pos_dict)
        basket_components = self.cta_engine.main_engine.get_basket_components(self.vt_symbol)
        print("basket_components:")
        print(basket_components)
        while True:
            basket_market_value = 0
            for comp in basket_components:
                if comp.exchange != etf_contract.exchange:
                    continue
                tick: TickData = self.cta_engine.main_engine.get_tick(comp.vt_symbol)
                if not tick:
                    continue
                pos = abs(leg_pos.get(comp.vt_symbol, 0))
                pos = min(pos, comp.share)
                if pos > 0:
                    basket_market_value += tick.useable_price * pos
                    leg_pos[comp.vt_symbol] = abs(leg_pos[comp.vt_symbol]) - pos

            if basket_market_value > 0.95 * stand_basket_market_value:
                basket_pos += 1
            else:
                basket_pos += basket_market_value / stand_basket_market_value
                break
        print(f"篮子持仓: basket_pos：  {basket_pos} ,basket_market_value：  {basket_market_value}, stand_basket_market_value： {stand_basket_market_value}")
        return basket_pos

    def buy_sell_with_target(
            self,
            limit_price: float,
            target_volume: float,
            per_order_max: float,
            signal_price: float = None,
            stop: bool = False,
            lock: bool = False,
            net: bool = False
    ):
        """
        根据目标仓位和每批次下单量etf下单
        :param target_volume: 目标仓位, 是1，2，。。个单位，不是具体的数量
        :param per_order_max: 分批下单本次最多的数量
        :return:
        """
        vol_gap = target_volume * self.pre_ss_vol - self.pos.get(self.vt_symbol, 0)   # 仓位缺口
        print(f"buy_sell_with_target: {vol_gap} , {target_volume} , {self.pre_ss_vol}, {self.pos.get(self.vt_symbol, 0)}")
        this_vol = min(abs(vol_gap), per_order_max)
        if vol_gap > 0 and this_vol < 100:
            this_vol = 100
        if vol_gap > 0:
            self.last_order_time = time.time()
            return self.buy(
                limit_price=limit_price,
                volume=this_vol,
                signal_price=signal_price,
                stop=stop,
                lock=lock,
                net=net
            )
        elif vol_gap < 0:
            self.last_order_time = time.time()
            return self.sell(
                limit_price=limit_price,
                volume=this_vol,
                signal_price=signal_price,
                stop=stop,
                lock=lock,
                net=net
            )
        return ""

    def purchase(self, volume):
        """
        申购
        """
        self.last_order_time = time.time()
        return self.send_order(price=0, volume=volume, direction=Direction.PURCHASE, offset=Offset.NONE,
                               lock=False)

    def redemption(self, volume):
        """
        赎回
        """
        self.last_order_time = time.time()
        return self.send_order(price=0, volume=volume, direction=Direction.REDEMPTION, offset=Offset.NONE,
                               lock=False)

    def set_basket_target(self, target_volume):
        """
        买卖篮子设置篮子目标仓位，系统会根据目标仓位进行计算买卖逻辑. 如果是卖出，应该设置target_volume=0
        :param target_volume:
        :return:
        """
        self.target_basket_pos = target_volume
        self.calc_basket_pos()
        order_requests = []
        etf_contract = self.cta_engine.main_engine.get_contract(self.vt_symbol)
        if target_volume == self.basket_pos:
            return
        for comp in self.cta_engine.main_engine.get_basket_components(self.vt_symbol):
            comp_symbol = comp.vt_symbol
            if not self.cta_engine.main_engine.is_basket_component_req_bs(comp, etf_contract):
                continue
            contract: ContractData = self.cta_engine.main_engine.get_contract(comp_symbol)
            if not contract:
                continue
            comp_to_bs = target_volume * comp.share - self.pos.get(comp_symbol, 0)
            tick: TickData = self.cta_engine.main_engine.get_tick(comp_symbol)
            if comp_to_bs > 0:
                order_requests.append(OrderRequest(
                    direction=Direction.LONG,
                    offset=Offset.OPEN,
                    price=tick.ask_price_5,
                    volume=comp_to_bs,
                    signal_price=None,
                    symbol=contract.symbol,
                    exchange=contract.exchange,
                    gateway_name=contract.gateway_name,
                    type=OrderType.BestOrLimit)
                )
            elif comp_to_bs < 0:
                real_pos = self.pos.get(comp_symbol, 0)
                if abs(comp_to_bs) < real_pos:
                    comp_to_bs = real_pos
                if comp_to_bs == 0:
                    continue
                order_requests.append(OrderRequest(
                    direction=Direction.SHORT,
                    offset=Offset.CLOSE,
                    price=tick.bid_price_5,
                    volume=abs(comp_to_bs),
                    signal_price=None,
                    symbol=contract.symbol,
                    exchange=contract.exchange,
                    gateway_name=contract.gateway_name,
                    type=OrderType.BestOrLimit)
                )
        order_ids = self.cta_engine.send_order_many(self, order_requests)
        self.last_order_time = time.time()
        self.active_orderids.update(order_ids)
        return order_ids

    def check_basket_req_bs(self, target_volume: float):
        """
        检查是否还有篮子股需要买卖
        :param target_volume:
        :return: True 为需要买卖， False为不需要买卖
        """
        etf_contract = self.cta_engine.main_engine.get_contract(self.vt_symbol)
        for comp in self.cta_engine.main_engine.get_basket_components(self.vt_symbol):
            comp_symbol = comp.vt_symbol
            if not self.cta_engine.main_engine.is_basket_component_req_bs(comp, etf_contract):
                continue
            contract: ContractData = self.cta_engine.main_engine.get_contract(comp_symbol)
            if not contract:
                continue
            comp_to_bs = target_volume * comp.share - self.pos.get(comp_symbol, 0)
            comp_to_bs = round_to(comp_to_bs, target=contract.pricetick)
            if abs(comp_to_bs) >= 100:
                # 只要一个需要买卖，就返回True
                return True
        # 篮子处理完成，已经不需要买卖
        return False

    def get_data(self):
        """
        Get strategy data.
        """
        data = super().get_data()
        data['trade_basket'] = self.trade_basket
        return data

    def get_report_data(self):
        return ReportStrategy(
            name=self.strategy_name,
            strategy_type='CTA-篮子',
            trading=self.trading,
            symbols={'A': self.vt_symbol},
            positions={
                self.vt_symbol: self.pos[self.vt_symbol],
                'basket': self.basket_pos,
                'all': self.all_pos
            },
            targets={},
            statue='success',
            client=WORK_DIR,
        ).pformat()